import os
import secrets
import requests
from datetime import datetime, timedelta
from flask import Flask, request, redirect, render_template, jsonify, session, url_for
from rich.console import Console
from rich.table import Table
from rich.panel import Panel
from rich.text import Text
from rich import print
import sqlite3
from urllib.parse import urlparse
from functools import wraps

app = Flask(__name__)
app.secret_key = secrets.token_hex(32)  # 用于会话加密
console = Console()

# ByUsi API 配置
BYUSI_API_URL = "https://api.www.byusi.cn/user/api.php"

# 数据库初始化
def init_db():
    conn = sqlite3.connect('shortlinks.db')
    c = conn.cursor()
    
    # 短链表
    c.execute('''CREATE TABLE IF NOT EXISTS shortlinks
                 (id INTEGER PRIMARY KEY AUTOINCREMENT,
                  short_code TEXT UNIQUE,
                  original_url TEXT,
                  created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                  expires_at TIMESTAMP,
                  hits INTEGER DEFAULT 0,
                  creator_uid TEXT,
                  is_public BOOLEAN DEFAULT 1,
                  title TEXT,
                  description TEXT)''')
    
    # 用户会话表
    c.execute('''CREATE TABLE IF NOT EXISTS user_sessions
                 (session_id TEXT PRIMARY KEY,
                  user_data TEXT,
                  created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                  last_activity TIMESTAMP)''')
    
    conn.commit()
    conn.close()

init_db()

# 生成短码
def generate_short_code(length=6):
    return secrets.token_urlsafe(length)[:length]

# 验证URL格式
def is_valid_url(url):
    try:
        result = urlparse(url)
        return all([result.scheme, result.netloc])
    except:
        return False

# 登录装饰器
def login_required(f):
    @wraps(f)
    def decorated_function(*args, **kwargs):
        if 'user' not in session:
            return redirect(url_for('login', next=request.url))
        return f(*args, **kwargs)
    return decorated_function

# 管理员装饰器
def admin_required(f):
    @wraps(f)
    def decorated_function(*args, **kwargs):
        if 'user' not in session or not session['user'].get('is_admin', False):
            return render_template('403.html'), 403
        return f(*args, **kwargs)
    return decorated_function

# Rich 控制台输出美化
def print_shortlink_info(short_code, original_url, expires_at, hits, creator_uid):
    table = Table(title="短链信息", show_header=True, header_style="bold magenta")
    table.add_column("属性", style="dim")
    table.add_column("值")
    
    table.add_row("短码", short_code)
    table.add_row("原始URL", original_url)
    table.add_row("过期时间", expires_at if expires_at else "永不过期")
    table.add_row("点击次数", str(hits))
    table.add_row("创建者UID", creator_uid)
    
    console.print(table)

# ByUsi API 调用
def call_byusi_api(action, data):
    data['action'] = action
    try:
        response = requests.post(BYUSI_API_URL, data=data)
        return response.json()
    except requests.exceptions.RequestException as e:
        console.print(f"[red]ByUsi API 调用失败: {e}[/]")
        return {'status': 'error', 'message': 'API服务不可用'}

@app.route('/')
def index():
    # 获取公开短链
    conn = sqlite3.connect('shortlinks.db')
    c = conn.cursor()
    c.execute("""
        SELECT short_code, original_url, title, description, hits 
        FROM shortlinks 
        WHERE is_public = 1 AND (expires_at IS NULL OR expires_at > datetime('now'))
        ORDER BY hits DESC 
        LIMIT 10
    """)
    popular_links = c.fetchall()
    conn.close()
    
    return render_template('index.html', popular_links=popular_links, user=session.get('user'))

@app.route('/register', methods=['GET', 'POST'])
def register():
    if request.method == 'POST':
        username = request.form.get('username')
        email = request.form.get('email')
        password = request.form.get('password')
        
        result = call_byusi_api('register', {
            'username': username,
            'email': email,
            'password': password
        })
        
        if result.get('status') == 'success':
            session['user'] = result['data']
            return redirect(url_for('dashboard'))
        
        return render_template('register.html', error=result.get('message'))
    
    return render_template('register.html')

@app.route('/login', methods=['GET', 'POST'])
def login():
    if request.method == 'POST':
        identifier = request.form.get('identifier')
        password = request.form.get('password')
        
        result = call_byusi_api('login', {
            'identifier': identifier,
            'password': password
        })
        
        if result.get('status') == 'success':
            session['user'] = result['data']
            
            # 生成唯一的 session_id
            session_id = secrets.token_urlsafe(32)
            session['session_id'] = session_id
            
            # 存储会话到数据库
            conn = sqlite3.connect('shortlinks.db')
            c = conn.cursor()
            c.execute("""
                INSERT INTO user_sessions (session_id, user_data, last_activity)
                VALUES (?, ?, datetime('now'))
                ON CONFLICT(session_id) DO UPDATE SET
                last_activity = datetime('now')
            """, (session_id, str(result['data'])))  # 修改这里
            conn.commit()
            conn.close()
            
            next_url = request.args.get('next', url_for('dashboard'))
            return redirect(next_url)
        
        return render_template('login.html', error=result.get('message'))
    
    return render_template('login.html')

@app.route('/logout')
def logout():
    if 'session_id' in session:
        # 从数据库删除会话
        conn = sqlite3.connect('shortlinks.db')
        c = conn.cursor()
        c.execute("DELETE FROM user_sessions WHERE session_id=?", (session['session_id'],))  # 修改这里
        conn.commit()
        conn.close()
    
    # 清除 session 数据
    session.pop('session_id', None)
    session.pop('user', None)
    return redirect(url_for('index'))

@app.route('/dashboard')
@login_required
def dashboard():
    user_uid = session['user']['uid']
    
    conn = sqlite3.connect('shortlinks.db')
    c = conn.cursor()
    
    # 用户短链统计
    c.execute("""
        SELECT COUNT(*), SUM(hits) 
        FROM shortlinks 
        WHERE creator_uid = ?
    """, (user_uid,))
    stats = c.fetchone()
    total_links = stats[0] or 0
    total_hits = stats[1] or 0
    
    # 最近创建的短链
    c.execute("""
        SELECT short_code, original_url, created_at, hits 
        FROM shortlinks 
        WHERE creator_uid = ?
        ORDER BY created_at DESC 
        LIMIT 5
    """, (user_uid,))
    recent_links = c.fetchall()
    
    # 热门短链
    c.execute("""
        SELECT short_code, original_url, hits 
        FROM shortlinks 
        WHERE creator_uid = ?
        ORDER BY hits DESC 
        LIMIT 5
    """, (user_uid,))
    top_links = c.fetchall()
    
    conn.close()
    
    return render_template('dashboard.html', 
                         user=session['user'],
                         total_links=total_links,
                         total_hits=total_hits,
                         recent_links=recent_links,
                         top_links=top_links)

@app.route('/create', methods=['POST'])
@login_required
def create_shortlink():
    try:
        # 参数获取与验证
        original_url = request.form.get('url', '').strip()
        custom_code = request.form.get('custom_code', '').strip()
        expire_days = int(request.form.get('expire_days', 30))
        title = request.form.get('title', '').strip()
        description = request.form.get('description', '').strip()
        is_public = 1 if request.form.get('is_public') == 'on' else 0

        # 增强参数验证
        if not original_url:
            return jsonify({'error': 'URL不能为空'}), 400

        if not is_valid_url(original_url):
            return jsonify({'error': '无效的URL格式，请包含http://或https://'}), 400

        if custom_code:
            if not re.match(r'^[a-zA-Z0-9_-]{4,20}$', custom_code):
                return jsonify({'error': '自定义短码只能包含字母、数字、下划线和连字符，长度4-20'}), 400

        # 数据库操作
        conn = None
        try:
            conn = sqlite3.connect('shortlinks.db')
            c = conn.cursor()
            
            # 处理自定义短码
            short_code = custom_code
            if short_code:
                c.execute("SELECT short_code FROM shortlinks WHERE short_code=?", (short_code,))
                if c.fetchone():
                    return jsonify({'error': '该短码已被使用'}), 400
            else:
                # 生成随机短码并保证唯一性
                max_retries = 5
                for _ in range(max_retries):
                    short_code = generate_short_code()
                    c.execute("SELECT short_code FROM shortlinks WHERE short_code=?", (short_code,))
                    if not c.fetchone():
                        break
                else:
                    return jsonify({'error': '短码生成失败，请重试或使用自定义短码'}), 500

            # 设置过期时间
            expires_at = (datetime.now() + timedelta(days=expire_days)).strftime('%Y-%m-%d %H:%M:%S') if expire_days > 0 else None

            # 插入数据库（包含事务处理）
            c.execute("""
                INSERT INTO shortlinks 
                (short_code, original_url, expires_at, creator_uid, is_public, title, description) 
                VALUES (?, ?, ?, ?, ?, ?, ?)
            """, (
                short_code,
                original_url,
                expires_at,
                session['user']['uid'],
                is_public,
                title[:50],  # 限制标题长度
                description[:200]  # 限制描述长度
            ))
            
            conn.commit()

            # 构建返回结果
            short_url = f"{request.host_url}{short_code}"
            console.print(f"[green]成功创建短链: {short_url} -> {original_url}")
            
            return jsonify({
                'short_url': short_url,
                'stats_url': f"{request.host_url}stats/{short_code}",
                'expires_at': expires_at
            })

        except sqlite3.IntegrityError as e:
            conn.rollback()
            console.print(f"[red]数据库唯一性冲突: {e}")
            return jsonify({'error': '短码冲突，请尝试其他短码'}), 409
            
        except Exception as e:
            conn.rollback()
            console.print(f"[red]数据库操作异常: {e}")
            return jsonify({'error': '数据库操作失败'}), 500
            
        finally:
            if conn:
                conn.close()

    except ValueError:
        return jsonify({'error': '无效的过期时间参数'}), 400
        
    except Exception as e:
        console.print(f"[red]系统异常: {e}")
        return jsonify({'error': '服务器内部错误'}), 500

@app.route('/<short_code>')
def redirect_shortlink(short_code):
    conn = sqlite3.connect('shortlinks.db')
    c = conn.cursor()
    
    c.execute("""
        SELECT original_url, expires_at, is_public, creator_uid 
        FROM shortlinks 
        WHERE short_code=?
    """, (short_code,))
    result = c.fetchone()
    
    if not result:
        conn.close()
        return render_template('404.html'), 404
    
    original_url, expires_at, is_public, creator_uid = result
    
    # 检查是否过期
    if expires_at and datetime.strptime(expires_at, '%Y-%m-%d %H:%M:%S') < datetime.now():
        conn.close()
        return render_template('expired.html'), 410
    
    # 检查私有短链权限
    if not is_public and ('user' not in session or session['user']['uid'] != creator_uid):
        conn.close()
        return render_template('403.html'), 403
    
    # 更新点击次数
    c.execute("UPDATE shortlinks SET hits = hits + 1 WHERE short_code=?", (short_code,))
    conn.commit()
    conn.close()
    
    return redirect(original_url)

@app.route('/admin')
@admin_required
def admin_dashboard():
    conn = sqlite3.connect('shortlinks.db')
    c = conn.cursor()
    
    # 系统统计
    c.execute("SELECT COUNT(*) FROM shortlinks")
    total_links = c.fetchone()[0]
    
    c.execute("SELECT SUM(hits) FROM shortlinks")
    total_hits = c.fetchone()[0] or 0
    
    c.execute("SELECT COUNT(*) FROM user_sessions")
    active_sessions = c.fetchone()[0]
    
    # 最近短链
    c.execute("""
        SELECT s.short_code, s.original_url, s.created_at, s.hits, u.username 
        FROM shortlinks s
        LEFT JOIN user_sessions us ON s.creator_uid = json_extract(us.user_data, '$.uid')
        LEFT JOIN (
            SELECT json_extract(user_data, '$.uid') as uid, 
                   json_extract(user_data, '$.username') as username 
            FROM user_sessions
            GROUP BY uid
        ) u ON s.creator_uid = u.uid
        ORDER BY s.created_at DESC 
        LIMIT 10
    """)
    recent_links = c.fetchall()
    
    # 活跃用户
    c.execute("""
    SELECT json_extract(user_data, '$.username') as username,
           json_extract(user_data, '$.uid') as uid,
           COUNT(*) as link_count,
           SUM(s.hits) as total_hits
    FROM user_sessions us
    LEFT JOIN shortlinks s ON json_extract(us.user_data, '$.uid') = s.creator_uid
    GROUP BY uid
    ORDER BY total_hits DESC
    LIMIT 10
""")
    top_users = c.fetchall()
    
    conn.close()
    
    return render_template('admin.html', 
                         user=session['user'],
                         total_links=total_links,
                         total_hits=total_hits,
                         active_sessions=active_sessions,
                         recent_links=recent_links,
                         top_users=top_users)

@app.route('/api/stats')
def get_stats():
    conn = sqlite3.connect('shortlinks.db')
    c = conn.cursor()
    
    # 总短链数
    c.execute("SELECT COUNT(*) FROM shortlinks")
    total_links = c.fetchone()[0]
    
    # 总点击量
    c.execute("SELECT SUM(hits) FROM shortlinks")
    total_hits = c.fetchone()[0] or 0
    
    # 今日新增
    c.execute("SELECT COUNT(*) FROM shortlinks WHERE date(created_at) = date('now')")
    today_new = c.fetchone()[0]
    
    # 活跃用户数
    c.execute("SELECT COUNT(DISTINCT creator_uid) FROM shortlinks")
    active_users = c.fetchone()[0]
    
    conn.close()
    
    return jsonify({
        'total_links': total_links,
        'total_hits': total_hits,
        'today_new': today_new,
        'active_users': active_users
    })

def print_startup_banner():
    banner = Text(r"""
 ____        _   _     _      ____  _                _   _     _       _    
| __ ) _   _| | | |___(_)    / ___|| |__   ___  _ __| |_| |   (_)_ __ | | __
|  _ \| | | | | | / __| |____\___ \| '_ \ / _ \| '__| __| |   | | '_ \| |/ /
| |_) | |_| | |_| \__ \ |_____|__) | | | | (_) | |  | |_| |___| | | | |   < 
|____/ \__, |\___/|___/_|    |____/|_| |_|\___/|_|   \__|_____|_|_| |_|_|\_\
       |___/                                                                
    """, style="bold blue")
    
    info_panel = Panel.fit(
        Text.from_markup(f"""
[b]短链服务已启动[/b]
        
• 访问地址: [green]{request.host_url}[/]
• 用户面板: [green]{request.host_url}dashboard[/]
• 管理面板: [green]{request.host_url}admin[/]
• API端点: [green]{request.host_url}api/stats[/]
        
使用 [yellow]Ctrl+C[/] 停止服务
        """),
        title="服务信息",
        border_style="yellow"
    )
    
    console.print(banner)
    console.print(info_panel)

@app.before_request
def log_request():
    if request.path == '/favicon.ico':
        return
    console.print(f"[dim]{datetime.now().strftime('%H:%M:%S')}[/] - [blue]{request.method}[/] [green]{request.path}[/] from [yellow]{request.remote_addr}[/]")

if __name__ == '__main__':
    with app.test_request_context('/'):
        print_startup_banner()
    app.run(host='0.0.0.0', port=55262, debug=True)